import os
import neologdn
from sklearn.feature_extraction.text import CountVectorizer, TfidfTransformer
import MeCab
import re
import matplotlib.pyplot as plt
from sklearn.cluster import KMeans
from sklearn.decomposition import PCA


# 清理不需要的字符
def clean_text(data):
    data = neologdn.normalize(data)
    return data


# 加载文件夹，返回文件名
def txt_load(folder_path, a):
    with open(os.path.join(folder_path, a), mode='r', encoding='utf-8')as f:
        data = f.read()
    return data


# 形態素解析，去空格
def fenci(data):
    mecab = MeCab.Tagger()
    words = mecab.parse(data)
    output = words.split('\n')
    items = (re.split('[\t,]', line) for line in output)
    return items


# 名詞抽出，只对名词词性的单词建模
def noun_extract(items):
    return ' '.join(item[0]
                    for item in items
                    if (item[0] not in 'EOS' and
                        item[1] == '名詞' and item[2] == '一般'))


def main():
    corpus = []
    folder_path = './data'
    folder_list = os.listdir(folder_path)
    for a in folder_list:
        data = txt_load('./data', a)
        data = clean_text(data)
        items = fenci(data)
        words = noun_extract(items)
        corpus.append(words)
    print(corpus)

    # 単語文章行列を作成　tf
    vectorized = CountVectorizer()
    count = vectorized.fit_transform(corpus)
    # 単語文章行列を作成tf-idf
    transformer = TfidfTransformer()
    tfidf_matrix = transformer.fit_transform(count)
    matrix1 = tfidf_matrix.toarray()
    # 主成分分析
    pca = PCA(n_components=2)
    pca.fit(matrix1)
    new_matrix1 = pca.fit_transform(matrix1)

    # K- means 无监督学习
    clf = KMeans(n_clusters=5, max_iter=100)
    clf.fit(new_matrix1)
    # mae误差
    print("MAE:", clf.inertia_)
    # 预测标签
    print(clf.labels_)

    for i in range(5):
        labels = new_matrix1[clf.labels_ == i]
        plt.scatter(labels[:, 0], labels[:, 1])
    centers = clf.cluster_centers_
    plt.scatter(centers[:, 0], centers[:, 1], facecolors='none', edgecolors='black')
    plt.title("クラスタリング結果（個数5）")
    plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
    plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号
    plt.show()


if __name__ == '__main__':
    main()
